from typing import Annotated

from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool, InjectedToolCallId
from langgraph.prebuilt import InjectedState
from langgraph.types import Command

from agent.my_state import CustomState


@tool
def get_username(tool_call_id: Annotated[str, InjectedToolCallId],
                 config: RunnableConfig) -> Command:
    """从configurable中获取用户名username，并更新到AgentState中，以便生成祝福语句"""
    print('call get username')
    username = config['configurable'].get('username', 'zhangsan')
    # print(f'调用工具，传入的用户名是：{username}')
    content = f'成功地获取了当前用户的username为：{username}'
    # Command调用一个命令，它的参数就是具体要执行的命令以及它的参数
    return Command(update={
        # 更新状态中的用户名
        'username': username,
        # 更新一条工具执行后的消息：ToolMessage类型
        # 注意这里会调用add_messages，会合并原先State中的Message与这里的Message
        'messages': [
            ToolMessage(
                content=content,
                # 这个tool_call_id是前一条message（一般是AIMessage）中指定的
                # 所以需要从前一个消息中获取这个tool_call_id中取得
                # 新版langgraph提供了一个依赖注入功能，在这个tool方法中定义一个参数tool_call_id
                # 如果不适用依赖注入功能，则需要从state中找到最后一条AIMessage，并获取其中的tool_call_id
                tool_call_id=tool_call_id
            )
        ]
    })


@tool
# 这里的CustomState也是自动注入的
def greet_user(state: Annotated[CustomState, InjectedState]) -> str:
    """在获取用户username之后，生成祝福语句。如果用户要求生成祝福语，请使用此工具！"""
    # 从state中获取用户名
    print('call greet user')
    username = state.get('username', 'anonymous_user')
    return f'祝贺你！（这是工具生成的祝福语！）{username}！'
